import torch
from torchvision import transforms as T
from torch.utils.data import Dataset,DataLoader
import pandas as pd
from PIL import Image

class ImageDs(Dataset):
    def __init__(self):
        super().__init__()
        self.df = pd.read_csv('./meta.csv', encoding="gbk")
        self.transformer = T.Compose([
            T.Resize((100, 100), antialias=True),
            T.Grayscale(),
            T.ToTensor()
        ])
        self.label_map = ['Apple', 'Peach']

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 根据idx获取一行数据，再拆分出路径 和 标签
        row = self.df.iloc[idx]     # 注意这里是[] 不是小括号
        img_path = row["img_path"]
        label = row["label"]
        # 路径用于PIL读取图像然后进行transformer转换为tensor张量
        img = Image.open(img_path)
        if self.transformer:
            img = self.transformer(img)
        # 标签用于label_map映射为 0 和 1的张量
        label = torch.tensor(self.label_map.index(label))
        return img, label

# 实例化对象
ds = ImageDs()

# 打印长度
print(len(ds))

# 找到图像和标签看一下
img, label = ds[0]
print(img.shape, label)

# 随机批量加载
dl = DataLoader(ds, batch_size=3, shuffle=True)
print(next(enumerate(dl)))

